#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/11/9 14:40
# @Author  : CHEN Wang
# @Site    : 
# @File    : factor_exposure_related.py
# @Software: PyCharm

"""
脚本说明: 获取因子暴露相关数据API
"""
import copy
from quant_researcher.quant.datasource_fetch.factor_api import factor_info
from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import BARRA_STYLE_FACTOR, \
    BARRA_INDUSTRY_FACTOR, BARRA_FULL_FACTOR
from quant_researcher.quant.project_tool.celebrity import get_lst_condition
from quant_researcher.quant.project_tool.db_operator import db_conn, my_mysql
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.logger.my_logger import LOG

T_BARRA_EXPOSURE = 'stk_barra_exposure'
T_SW_INDUSTRY_EXPOSURE = 'stk_fac_ind_sw_ind1'
T_ZZ_INDUSTRY_EXPOSURE = 'stk_fac_ind_zz_800_ind1'
EARLIEST_START_DATE = '1900-01-01'


def get_barra_factor_exposure(stock_code=None, start_date=None, end_date=None,
                              only_latest=False, factor_type='style'):
    """
    获取股票的barra因子暴露数据

    :param str or list or None stock_code: 股票代码，支持单个(str)多个(list)所有(None)
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param factor_type: style-风格因子 industry-行业因子 all-全因子
    :return: pd.DataFrame
    """
    if factor_type == 'style':
        select = copy.deepcopy(BARRA_STYLE_FACTOR)
    elif factor_type == 'industry':
        select = copy.deepcopy(BARRA_INDUSTRY_FACTOR)
    elif factor_type == 'all':
        select = copy.deepcopy(BARRA_FULL_FACTOR)
    else:
        raise NotImplementedError('目前只有风格因子、行业因子和全因子')
    select = ['stockcode as stock_code', 'tradedate as end_date'] + select

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    where = []

    if isinstance(stock_code, str) or isinstance(stock_code, list):
        if isinstance(stock_code, str):
            stock_code = [stock_code]
        stock_where_condition = get_lst_condition(stock_code, 'stockcode')
        where.append(stock_where_condition)
    elif stock_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    conn = db_conn.get_tk_factors_conn()

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=T_BARRA_EXPOSURE, select='max(tradedate)',
                                       where=[f"tradedate <= '{end_date}'", stock_where_condition],
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{T_BARRA_EXPOSURE}表为空，请检查")
            return
        where.append(f"tradedate = '{latest_date}'")
    else:
        where.append(f"tradedate >= '{start_date}'")
        where.append(f"tradedate <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=T_BARRA_EXPOSURE, select=select, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{factor_type}在{start_date}到{end_date}之间的数据为空，请检查")
        return

    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_stock_factor_exposure(factor, stock_code=None, start_date=None, end_date=None, only_latest=False):
    """
    获取股票因子暴露数据, todo 获取数据较慢，需要提速

    :param str factor: 因子英文名
    :param str or list or None stock_code: 股票代码，支持单个（str）、多个（list）、所有（None）
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame
    """
    print(f'获取因子：{factor}的数据')
    conn = db_conn.get_tk_factors_conn(by_sqlalchemy=True)
    factor_table_name, _ = factor_info.get_factor_table_name(factor)

    select = ['tradedate as end_date', 'stockcode as stock_code', f'value as {factor}']

    where = []

    if isinstance(stock_code, str) or isinstance(stock_code, list):
        if isinstance(stock_code, str):
            stock_code = [stock_code]
        stock_where_condition = get_lst_condition(stock_code, 'stockcode')
        where.append(stock_where_condition)
    elif stock_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=factor_table_name, select='max(tradedate)',
                                       where=[f"tradedate <= '{end_date}'", stock_where_condition],
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{factor_table_name}表为空，请检查")
            return
        where.append(f"tradedate = '{latest_date}'")
    else:
        where.append(f"tradedate >= '{start_date}'")
        where.append(f"tradedate <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=factor_table_name, select=select, where=where, conn=conn)
    # conn.close()

    if df.empty:
        LOG.error(f"{factor}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_industry_composite_factor_exposure(factor, industry_code=None, start_date=None, end_date=None, only_latest=False, method='size_weighted'):
    """
    获取行业合成暴露数据, todo 获取数据较慢，需要提速

    :param str factor: 因子英文名
    :param str or list or None industry_code: 股票代码，支持单个（str）、多个（list）、所有（None）
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :param method: 加权方式，默认size_weighted
    :return: pd.DataFrame
    """
    print(f'获取因子：{factor}的数据')
    conn = db_conn.get_tk_factors_conn(by_sqlalchemy=True)
    factor_table_name, _ = factor_info.get_factor_table_name(factor, asset_type='industry')

    select = ['tradedate as end_date', 'indexcode as industry_code', f'value as {factor}']

    where = []
    where.append(f"method = '{method}'")

    if isinstance(industry_code, str) or isinstance(industry_code, list):
        if isinstance(industry_code, str):
            industry_code = [industry_code]
        stock_where_condition = get_lst_condition(industry_code, 'indexcode')
        where.append(stock_where_condition)
    elif industry_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前industry_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=factor_table_name, select='max(tradedate)',
                                       where=[f"tradedate <= '{end_date}'", stock_where_condition],
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{factor_table_name}表为空，请检查")
            return
        where.append(f"tradedate = '{latest_date}'")
    else:
        where.append(f"tradedate >= '{start_date}'")
        where.append(f"tradedate <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=factor_table_name, select=select, where=where, conn=conn)
    # conn.close()

    if df.empty:
        LOG.error(f"{factor}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


def get_industry_factor_exposure(industry_type='SW', stock_code=None,
                                 start_date=None, end_date=None, only_latest=False):
    """
    获取行业因子暴露数据

    :param str industry_type: 行业因子类型，支持'SW'-申万, 'ZZ'-中证
    :param str or list or None stock_code: 股票代码，支持单个（str）、多个（list）、所有（None）
    :param str or int or datetime.date or datetime.datetime or int start_date: 开始时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param str or int or datetime.date or datetime.datetime or int end_date: 结束时间，支持带杠字符串、不带杠字符串、datetime、date、int格式
    :param bool only_latest: True-只获取最新的一天数据，False-获取一个区间段的数据
    :return: pd.DataFrame
    """
    conn = db_conn.get_tk_factors_conn()
    if industry_type == 'SW':
        table_name = T_SW_INDUSTRY_EXPOSURE
        indu_name = 'sw_ind1'
    elif industry_type == 'ZZ':
        table_name = T_ZZ_INDUSTRY_EXPOSURE
        indu_name = 'zz_800_ind1'
    else:
        raise NotImplementedError('industry_type目前只支持SW、ZZ')

    select = ['tradedate as end_date', 'stockcode as stock_code', indu_name]

    where = []

    if isinstance(stock_code, str) or isinstance(stock_code, list):
        if isinstance(stock_code, str):
            stock_code = [stock_code]
        stock_where_condition = get_lst_condition(stock_code, 'stockcode')
        where.append(stock_where_condition)
    elif stock_code is None:
        stock_where_condition = '1 = 1'
    else:
        LOG.error('目前stock_code只支持str、list或者None，请按要求输入')
        raise NotImplementedError

    if start_date is None:
        start_date = EARLIEST_START_DATE
    else:
        start_date = time_tool.format_date_str(start_date)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    else:
        end_date = time_tool.format_date_str(end_date)

    if only_latest:
        latest_date = my_mysql.read_v2(sfrom=table_name, select='max(tradedate)',
                                       where=[f"tradedate <= '{end_date}'", stock_where_condition],
                                       conn=conn).iloc[0, 0]
        if latest_date is None:
            LOG.error(f"{table_name}表为空，请检查")
            return
        where.append(f"tradedate = '{latest_date}'")
    else:
        where.append(f"tradedate >= '{start_date}'")
        where.append(f"tradedate <= '{end_date}'")

    df = my_mysql.read_v2(sfrom=table_name, select=select, where=where, conn=conn)
    conn.close()

    if df.empty:
        LOG.error(f"{industry_type}在{start_date}到{end_date}之间的数据为空，请检查")
        return
    df['end_date'] = df['end_date'].apply(time_tool.format_date_str)

    return df


"""
之前的API函数，目前都替换成最新的API函数，以下函数已弃用
"""


# def get_barra_factor_exposure(stock_code=None, start_date=None, end_date=None, factor_type='style'):
#     """
#     获取barra因子在指定时间区间的暴露数据
#
#     :param str start_date: 开始日期, 如："2020-01-08"，如果只获取一天，设为None
#     :param str end_date: 截止日期, 如："2020-01-08"，如果start_date和end_date都为None，获取所有数据
#     :param factor_type: style: 风格因子 industry: 行业因子 all: 全因子
#     :param stock_code: 股票代码，支持单个(str)多个(list)所有(None)
#     :return:
#     """
#     conn = db_conn.get_tk_factors_conn()
#     if factor_type == 'style':
#         factor_list = BARRA_STYLE_FACTOR
#     elif factor_type == 'industry':
#         factor_list = BARRA_INDUSTRY_FACTOR
#     elif factor_type == 'all':
#         factor_list = BARRA_STYLE_FACTOR+BARRA_INDUSTRY_FACTOR
#     else:
#         raise ValueError('目前仅支持style，industry，all三种类型的因子获取')
#     where = []
#     if end_date is not None and start_date is None:
#         end_date = pd.read_sql(f"select max(tradedate) "
#                                f"from {T_BARRA_EXPOSURE} "
#                                f"where tradedate <='{end_date}'", conn).iloc[0, 0]
#         start_date = end_date
#     if end_date is None and start_date is None:
#         end_date = time_tool.get_yesterday(marker='with_n_dash')
#         start_date = '2011-01-01'
#
#     if start_date is not None:
#         where.append(f"tradedate >= '{start_date}'")
#     if end_date is not None:
#         where.append(f"tradedate <= '{end_date}' ")
#
#     if isinstance(stock_code, str):
#         where.append(f'stockcode=\'{stock_code}\'')
#     elif isinstance(stock_code, list):
#         where.append(get_lst_condition(stock_code, 'stock_code'))
#     elif stock_code is None:
#         pass
#     else:
#         raise NotImplementedError('目前stock_code只支持str、list或者None，请按要求输入')
#
#     select = ['stockcode']+factor_list
#     exposure = my_mysql.read_v2(
#         select=select, where=where, sfrom=T_BARRA_EXPOSURE,
#         conn=conn)
#     conn.close()
#     return exposure


# def get_stock_factor_exposure(factor_name_en, stock_pool=None, start=None, end=None):
#     """
#     获取股票因子暴露数据, todo 获取数据较慢，需要提速
#
#     :param str factor_name_en: 因子英文名
#     :param str stock_pool: 指定股票池， 默认为None, 则获取所有股票该因子数据
#     :param str start: 起始时间， '2020-01-01', 默认为None
#     :param str end: 结束时间， '2020-05-01', 默认为None, 只输入end则获取这一天的截面暴露;
#                     若start, end皆为None，则获取该因子所有暴露数据
#     :return:
#     """
#
#     print(f'获取因子：{factor_name_en}的数据')
#     conn = db_conn.get_tk_factors_conn()
#     factor_table_name, _ = factor_info.get_factor_table_name(factor_name_en)
#     if end is None and start is None:
#         end = time_tool.get_yesterday(marker='with_n_dash')
#         start = '2000-01-01'
#     elif end is not None and start is None:
#         start = end
#     where_condition = f"tradedate >= '{start}' and tradedate <= '{end}' "
#
#     if stock_pool is not None:
#         where_condition += f"and stockcode in {tuple(stock_pool + [''])}"
#
#     factor_data = pd.read_sql(f"select tradedate, stockcode, value as {factor_name_en} "
#                               f"from {factor_table_name} "
#                               f"where {where_condition} "
#                               f"order by tradedate ", conn)
#     conn.close()
#     return factor_data


# def get_industry_factor_exposure(industry_type='SW', stock_pool=None, start=None, end=None):
#     """
#     获取行业因子暴露数据
#
#     :param str industry_type: 行业因子类型，支持'SW'：申万, 'ZZ'：中证
#     :param str stock_pool: 指定股票池， 默认为None, 则获取所有股票行业因子数据
#     :param str start: 起始时间， '2020-01-01', 默认为None
#     :param str end: 结束时间， '2020-05-01', 默认为None, 只输入end则获取这一天的截面行业暴露;
#                     若start, end皆为None，则获取所有暴露数据
#     :return:
#     """
#
#     conn = db_conn.get_tk_factors_conn()
#     if industry_type == 'SW':
#         table_name = 'stk_fac_ind_sw_ind1'
#         indu_name = 'sw_ind1'
#     elif industry_type == 'ZZ':
#         table_name = 'stk_fac_ind_zz_800_ind1'
#         indu_name = 'zz_800_ind1'
#     else:
#         raise NotImplementedError('industry_type目前只支持SW、ZZ')
#     if end is None and start is None:
#         end = time_tool.get_yesterday(marker='with_n_dash')
#         start = '2000-01-01'
#     elif end is not None and start is None:
#         start = end
#     where_condition = f"tradedate >= '{start}' and tradedate <= '{end}' "
#
#     if stock_pool is not None:
#         where_condition += f"and stockcode in {tuple(stock_pool + [''])}"
#
#     ind = pd.read_sql(f"select tradedate, stockcode, {indu_name} from {table_name} "
#                       f"where {where_condition} ", conn)
#     conn.close()
#     return ind


if __name__ == '__main__':
    # aaa = get_barra_factor_exposure(None, '2020-01-01', '2020-02-01', False, 'all')
    # bbb = get_stock_factor_exposure('cetop', ['000001', '000002'], '2020-01-01', '2020-03-20', False)
    # ccc = get_industry_factor_exposure('SW', ['000001', '000002'], '2020-01-01', '2020-02-01', False)
    ccc = get_industry_composite_factor_exposure('e_grow_f1', ['000928', '801010'], '2020-01-01', '2020-02-01', False)
